import os
import numpy as np
import torch
from torch.utils.data import DataLoader


class CustomDataset(torch.utils.data.Dataset):
    def __init__(self,data_path):
        data = np.load(data_path, allow_pickle=True)
        self.images = torch.tensor(data['images'], dtype=torch.float32)
        self.labels = torch.tensor(data['labels'], dtype=torch.int64)

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, idx):
        image = self.images[idx]
        label = self.labels[idx]
        return image, label

def load_client_data(is_train=True,datafile='data1'):
    data_type = 'train' if is_train else 'test'
    data_path = os.path.join(f'./data/{data_type}/{datafile}')
    dataset = CustomDataset(data_path)

    return dataset

